package cn.tworice.chatbot.core.client.baidu;

import cn.tworice.chatbot.config.ChatBotProperties;
import cn.tworice.chatbot.core.bo.ChatMessage;
import cn.tworice.chatbot.core.client.ChatBotClient;
import com.baidubce.qianfan.Qianfan;
import com.baidubce.qianfan.core.builder.ChatBuilder;
import com.baidubce.qianfan.model.chat.Message;

import javax.annotation.Resource;
import java.util.List;

/**
 * 百度的ERNIE大模型
 */
public class BaiduClient implements ChatBotClient {

    @Resource
    private ChatBotProperties chatBotProperties;

    @Override
    public String chat(String question, List<ChatMessage> msgList) {
        Qianfan qianfan = new Qianfan(chatBotProperties.getApiKey(), chatBotProperties.getSecretKey());
        ChatBuilder model = qianfan.chatCompletion()
                .model(chatBotProperties.getModel());
        for (ChatMessage msg : msgList) {
            model.addMessage(new Message().setRole(msg.getRole()).setContent(msg.getContent()).setName(msg.getName()));
        }
        return model.execute().getResult();
    }

    @Override
    public String getRoleUser() {
        return "user";
    }

    @Override
    public String getRoleAssistant() {
        return "assistant";
    }
}
